import numpy as np
from typing import Union, List, Any
from numpy import ndarray, dtype


class Utils:
    @staticmethod
    def binary(x: np.ndarray) -> np.ndarray:
        """
            f(x) = 1 or 0
        """
        return (x > 0).astype(float)

    @staticmethod
    def trivalent(x: np.ndarray) -> np.ndarray:
        """
            f(x) = 1 or 0 or -1
        """
        return np.where(
            x > 0, np.array(1.0),
            np.where(
                x < 0, np.array(-1.0), np.array(0.0))
        )

    @staticmethod
    def sigmoid(x: np.ndarray, _c: float = 1.0) -> np.ndarray:
        """sigmoid transformation function with configurable slope"""
        return 1 / (1 + np.exp(-_c * x))

    @staticmethod
    def tanh(x: np.ndarray, _c: float = 1.0) -> np.ndarray:
        return np.tanh(_c * x)

    funcs_registry = {
        'binary': binary,
        'trivalent': trivalent,
        'sigmoid': sigmoid,
        'tanh': tanh,
    }

    @staticmethod
    def print_trajectory(
            _i: int,
            current_state: np.ndarray,
            weighted_state: np.ndarray,
            activated_state: np.ndarray
    ):
        print("Current iteration: ", _i)
        print("current state: ", current_state)
        print("weighted state: ", weighted_state)
        print("After activate: ", activated_state)


class CognitiveMap:
    """
    NumPy Implementation of Cognitive Map Reasoning System
    Core iteration Formula (6): x_i(t+1) = f(∑_{j≠i} w_ij * x_j(t))
    Enhanced with consolidated data validation functions
    """

    def __init__(
            self,
            num_nodes: int,
            node_names: List[str],
            init_state: np.ndarray,
            init_weight: Union[np.ndarray, np.ndarray],
            act_fn_name: str,
            max_iterations: int,
            convergence_threshold: float = 1e-6,
    ):
        # load init data
        self.num_nodes = num_nodes
        self.node_names = node_names
        self.init_state = init_state
        self.init_weight = init_weight
        self.utils = Utils()
        self.activate_function = self.utils.funcs_registry[act_fn_name]

        # control parameters
        self.max_iterations = max_iterations
        self.convergence_threshold = convergence_threshold


    def model1_iteration(self):
        # Init input parameters
        current_state = self.init_state
        current_weight = self.init_weight
        state_history = [current_state.copy()]
        diff_history = []

        for i in range(self.max_iterations):
            # 计算 权重和, 然后用tanh激活函数
            # Calculate the sum of weight and then, active with tanh
            state_update = np.dot(current_state, current_weight)
            new_state = self.activate_function(state_update)

            state_diff = np.linalg.norm(new_state - current_state)
            diff_history.append(np.max(state_diff))
            state_history.append(new_state.copy())
            current_state =  new_state    # update input state

        return np.array(state_history), np.array(diff_history)


    def model3_iteration(
            self,
            _gamma,
            _lambda
    ) -> tuple[ndarray[Any, dtype[Any]], ndarray[Any, dtype[Any]], ndarray[Any, dtype[Any]]]:
        """
            Perform cognitive map reasoning process with consolidated validation
            F1: update state
                x_i(t+1) = tanh( sum_{j=1}^{n} w_{ij}(t) * x_j(t) + γ * x_i(t) )
            F2: update weight
                w_{ij}(t+1) = tanh( λ * w_{ij}(t) + (1-λ) * (x_j(t+1) - x_i(t+1)) )
        """

        # Init input parameters
        current_state = self.init_state
        current_weight = self.init_weight

        # Append history file
        state_history = [current_state.copy()]
        weight_history = [current_weight.copy()]
        diff_history = []

        for i in range(self.max_iterations):
            # 计算 权重和 + 自反馈项, 然后用tanh激活函数
            state_update = np.dot(current_state, current_weight) + _gamma * current_state
            new_state = self.activate_function(state_update)  #

            print("input_state", "\n", current_state)
            print("state_update", "\n", state_update)
            print("new_state", "\n", new_state)
            # 计算状态差异矩阵
            state_diff_matrix = np.subtract.outer(new_state, new_state)  # x_j - x_i

            # 根据更新之后的状态 更新权重矩阵
            weight_update = _lambda * current_weight + (1 - _lambda) * state_diff_matrix
            new_weight = self.activate_function(weight_update)

            print("input_weight", "\n", current_weight)
            print("diff_matrix", "\n", state_diff_matrix)
            print("weight_update", "\n", weight_update)
            print("new_state", "\n", new_state)

            # save output state
            state_diff = np.linalg.norm(new_state - current_state)
            diff_history.append(np.max(state_diff))

            state_history.append(new_state.copy())
            weight_history.append(new_weight.copy())

            current_state =  new_state    # update input state
            current_weight = new_weight

        return np.array(state_history), np.array(weight_history), np.array(diff_history)


if __name__ == "__main__":
    n_nodes = 5
    init_vector = np.array([1.0, 0, 0, 0, 0])
    init_weights = np.array(
        [
            [0, 0, -1, 0, 1],
            [0, 0, 0, -1, 0],
            [0, -1, 0, 0, -1],
            [-1, 1, 0, 0, 0],
            [0, 0, 0, 1, 1],
        ], dtype=np.float32)

    test1 = CognitiveMap(
        num_nodes=n_nodes,
        node_names=['1', '2', '3', '4', '5'],
        init_state=init_vector,
        init_weight=init_weights,
        act_fn_name='tanh',
        max_iterations= 10
    )

    a, b, c = test1.model3_iteration(0.5, 0.5)
    print(a, b, c)

    #r_history, d_history = test1.iteration()
    #print(r_history)
    #print(d_history)
    # l=0.5
    # g=0.5
    # v, w = test1.reasoning_model2(init_vector, init_weights, g, l)


# 1.20389753 1.39802546 1.15321794 0.74420728 0.3612557
# 0.17171826 0.08829389 0.04525667 0.02286045 0.01147083

